# -*- coding: utf-8 -*-
import os
import sys

current_path = os.path.dirname(__file__)
os.chdir(current_path)


# data_root_path = r'E:\shenhaitaoPyCode\leidiQA_web\dataset'


def create_data_list(data_root_path):
    '''with open(data_root_path + 'test_list.txt', 'w') as f:
        pass
    with open(data_root_path + 'train_list.txt', 'w') as f:
        pass'''

    with open(os.path.join(data_root_path, 'dict_txt.txt'), 'r', encoding='utf-8') as f_data:
        dict_txt = eval(f_data.readlines()[0])

    with open(os.path.join(data_root_path, 'train.txt'), 'r', encoding='utf-8') as f_data:
        lines = f_data.readlines()
    i = 0
    for line in lines:
        # title = line.split('_!_')[-1].replace('\n', '')
        # l = line.split('_!_')[1]
        # title = line.split('\t')[-1]   #ty

        # title = line.split('\t')[-1]
        # l = line.split('\t')[0]
        title = line.split('|')[-1]
        l = line.split('|')[0]
        labs = ""
        # if i % 10 == 0:
        #     with open(os.path.join(data_root_path, 'test_list.txt'), 'a', encoding='utf-8') as f_test:
        #         for s in title:
        #             # print(s)
        #             if s != " ":
        #                 lab = str(dict_txt[s])
        #                 labs = labs + lab + ','
        #         labs = labs[:-1]
        #         labs = labs + '\t' + l + '\n'
        #         f_test.write(labs)
        # else:
        #     with open(os.path.join(data_root_path, 'train_list.txt'), 'a', encoding='utf-8') as f_train:
        #         for s in title:
        #             if s != " ":
        #                 lab = str(dict_txt[s])
        #                 labs = labs + lab + ','
        #         labs = labs[:-1]
        #         labs = labs + '\t' + l + '\n'
        #         f_train.write(labs)
        with open(os.path.join(data_root_path, 'train_list.txt'), 'a', encoding='utf-8') as f_train:
            for s in title:
                if s != " ":
                    lab = str(dict_txt[s])
                    labs = labs + lab + ','
            labs = labs[:-1]
            labs = labs + '\t' + l + '\n'
            f_train.write(labs)
        i += 1
    print("数据列表生成完成！")


# 把数据生成一个字典
def create_dict(data_path, dict_path):
    dict_set = set()
    # 读取数据
    with open(data_path, 'r', encoding='utf-8') as f:
        lines = f.readlines()
    # 把数据生成一个元组
    for line in lines:
        title = line.split('|')[-1]
        # title = line.split('\t')[-1]
        for s in title:
            if s != " ":
                dict_set.add(s)
    # 把元组转换成字典，一个字对应一个数字
    dict_list = []
    i = 0
    for s in dict_set:
        dict_list.append([s, i])
        i += 1
    # 添加未知字符
    dict_txt = dict(dict_list)
    end_dict = {"<unk>": i}
    dict_txt.update(end_dict)
    # 把这些字典保存到本地
    with open(dict_path, 'w', encoding='utf-8') as f:
        f.write(str(dict_txt))

    print("数据字典生成成功")


data_root_path = "../dataset"


def start_prepare_dataset():
    # 把生产的数据列表都放在自己的总类别文件中
    # data_root_path = "E:/shenhaitaoPyCode/leidiQA_web/dataset"
    data_path = os.path.join(data_root_path, 'train.txt')
    dict_path = os.path.join(data_root_path, 'dict_txt.txt')

    # 创建数据字典
    create_dict(data_path, dict_path)
    # 创建数据列表
    create_data_list(data_root_path)


def save_dataset(questions, answer):
    if os.path.exists(data_root_path + '/dict_txt.txt'):  # 如果文件存在
        # 删除文件，可使用以下两种方法。
        os.remove(data_root_path + '/dict_txt.txt')
        # os.remove(data_root_path + '/test_list.txt')
        os.remove(data_root_path + '/train_list.txt')
    else:
        print('no such file:', "dict_txt.txt,train_list.txt")
    with open(os.path.join(data_root_path, 'train.txt'), 'r', encoding='utf-8') as f_data:
        lines = f_data.readlines()
        lastLine = lines[len(lines) - 1]
        lastLine = lastLine.split("|")[0]
        split_questions = questions.split("|")
    with open(os.path.join(data_root_path, 'train.txt'), 'a', encoding='utf-8') as f_data:
        for question in split_questions:
            line_ = int(lastLine) + 1
            add_question = str(line_) + "|" + question.replace("\n", "")
            f_data.write(add_question.strip() + "\n")
    with open(os.path.join(data_root_path, 'FAQ.txt'), 'r', encoding='utf-8') as f_data:
        lines = f_data.readlines()
        lastLine = lines[len(lines) - 1]
        lastLine = lastLine.split("|")[0]
    with open(os.path.join(data_root_path, 'FAQ.txt'), 'a', encoding='utf-8') as f_data:
        line_ = int(lastLine) + 1
        add_answer = str(line_) + "|" + answer.replace("\n", "")
        f_data.write(add_answer.strip() + "\n")
    print('储存完成，开始准备数据')
    start_prepare_dataset()


if __name__ == '__main__':
    # q = sys.argv[1]
    # a = sys.argv[2]
    # print("问题：", q)
    # print("答案：", a)
    # save_dataset(q, a)
    # print('数据添加成功')

    start_prepare_dataset()